{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#coding=gbk\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "#定义常量\n",
    "rnn_unit=10       #hidden layer units\n",
    "input_size=7\n",
    "output_size=1\n",
    "PATH = \"model\"\n",
    "lr=0.0006         #学习率\n",
    "#——————————————————导入数据——————————————————————\n",
    "f=open('dataset/dataset_2.csv')\n",
    "df=pd.read_csv(f)     #读入股票数据\n",
    "data=df.iloc[:,2:10].values  #取第3-10列\n",
    "\n",
    "\n",
    "#获取训练集\n",
    "def get_train_data(batch_size=60,time_step=20,train_begin=1,train_end=5800):\n",
    "    batch_index=[]\n",
    "    data_train=data[train_begin:train_end]\n",
    "    normalized_train_data=(data_train-np.mean(data_train,axis=0))/np.std(data_train,axis=0)  #标准化\n",
    "    train_x,train_y=[],[]   #训练集 \n",
    "    for i in range(len(normalized_train_data)-time_step):\n",
    "       if i % batch_size==0:\n",
    "           batch_index.append(i)\n",
    "       x=normalized_train_data[i:i+time_step,:7]\n",
    "       y=normalized_train_data[i:i+time_step,7,np.newaxis]\n",
    "       train_x.append(x.tolist())\n",
    "       train_y.append(y.tolist())\n",
    "    batch_index.append((len(normalized_train_data)-time_step))\n",
    "    return batch_index,train_x,train_y\n",
    "\n",
    "\n",
    "\n",
    "#获取测试集\n",
    "def get_test_data(time_step=20,test_begin=5800):\n",
    "    data_test=data[test_begin:]\n",
    "    mean=np.mean(data_test,axis=0)\n",
    "    std=np.std(data_test,axis=0)\n",
    "    normalized_test_data=(data_test-mean)/std  #标准化\n",
    "    size=(len(normalized_test_data)+time_step-1)//time_step  #有size个sample \n",
    "    test_x,test_y=[],[]  \n",
    "    for i in range(size-1):\n",
    "       x=normalized_test_data[i*time_step:(i+1)*time_step,:7]\n",
    "       y=normalized_test_data[i*time_step:(i+1)*time_step,7]\n",
    "       test_x.append(x.tolist())\n",
    "       test_y.extend(y)\n",
    "    test_x.append((normalized_test_data[(i+1)*time_step:,:7]).tolist())\n",
    "    test_y.extend((normalized_test_data[(i+1)*time_step:,7]).tolist())\n",
    "    return mean,std,test_x,test_y\n",
    "\n",
    "\n",
    "\n",
    "#——————————————————定义神经网络变量——————————————————\n",
    "#输入层、输出层权重、偏置\n",
    "\n",
    "weights={\n",
    "         'in':tf.Variable(tf.random_normal([input_size,rnn_unit])),\n",
    "         'out':tf.Variable(tf.random_normal([rnn_unit,1]))\n",
    "        }\n",
    "biases={\n",
    "        'in':tf.Variable(tf.constant(0.1,shape=[rnn_unit,])),\n",
    "        'out':tf.Variable(tf.constant(0.1,shape=[1,]))\n",
    "       }\n",
    "\n",
    "#——————————————————定义神经网络变量——————————————————\n",
    "def lstm(X):     \n",
    "    batch_size=tf.shape(X)[0]\n",
    "    time_step=tf.shape(X)[1]\n",
    "    w_in=weights['in']\n",
    "    b_in=biases['in']  \n",
    "    input=tf.reshape(X,[-1,input_size])  #需要将tensor转成2维进行计算，计算后的结果作为隐藏层的输入\n",
    "    input_rnn=tf.matmul(input,w_in)+b_in\n",
    "    input_rnn=tf.reshape(input_rnn,[-1,time_step,rnn_unit])  #将tensor转成3维，作为lstm cell的输入\n",
    "    cell=tf.nn.rnn_cell.BasicLSTMCell(rnn_unit)\n",
    "    init_state=cell.zero_state(batch_size,dtype=tf.float32)\n",
    "    output_rnn,final_states=tf.nn.dynamic_rnn(cell, input_rnn,initial_state=init_state, dtype=tf.float32)  #output_rnn是记录lstm每个输出节点的结果，final_states是最后一个cell的结果\n",
    "    output=tf.reshape(output_rnn,[-1,rnn_unit]) #作为输出层的输入\n",
    "    w_out=weights['out']\n",
    "    b_out=biases['out']\n",
    "    pred=tf.matmul(output,w_out)+b_out\n",
    "    return pred,final_states\n",
    "\n",
    "\n",
    "\n",
    "#——————————————————训练模型——————————————————\n",
    "def train_lstm(batch_size=80,time_step=15,train_begin=2000,train_end=5800):\n",
    "    X=tf.placeholder(tf.float32, shape=[None,time_step,input_size])\n",
    "    Y=tf.placeholder(tf.float32, shape=[None,time_step,output_size])\n",
    "    batch_index,train_x,train_y=get_train_data(batch_size,time_step,train_begin,train_end)\n",
    "    pred,_=lstm(X)\n",
    "    #损失函数\n",
    "    loss=tf.reduce_mean(tf.square(tf.reshape(pred,[-1])-tf.reshape(Y, [-1])))\n",
    "    train_op=tf.train.AdamOptimizer(lr).minimize(loss)\n",
    "    saver=tf.train.Saver(tf.global_variables(),max_to_keep=15)\n",
    "    # module_file = tf.train.latest_checkpoint()\n",
    "    with tf.Session() as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        ckpt = tf.train.get_checkpoint_state(PATH)\n",
    "        if ckpt and ckpt.model_checkpoint_path:\n",
    "            saver.restore(sess, ckpt.model_checkpoint_path)\n",
    "            print(\"Load!\")\n",
    "        else:\n",
    "            tf.global_variables_initializer().run()\n",
    "        # saver.restore(sess, module_file)\n",
    "        #重复训练10000次\n",
    "        for i in range(100):\n",
    "            for step in range(len(batch_index)-1):\n",
    "                _,loss_=sess.run([train_op,loss],feed_dict={X:train_x[batch_index[step]:batch_index[step+1]],Y:train_y[batch_index[step]:batch_index[step+1]]})\n",
    "            print(i,loss_)\n",
    "            if i % 200==0:\n",
    "                saver.save(sess, 'model/model.ckpt', global_step=i)\n",
    "                print(\"保存模型：\")\n",
    "\n",
    "#\n",
    "# train_lstm()\n",
    "\n",
    "\n",
    "#————————————————预测模型————————————————————\n",
    "def prediction(time_step=20):\n",
    "    X=tf.placeholder(tf.float32, shape=[None,time_step,input_size])\n",
    "    #Y=tf.placeholder(tf.float32, shape=[None,time_step,output_size])\n",
    "    mean,std,test_x,test_y=get_test_data(time_step)\n",
    "    pred,_=lstm(X)     \n",
    "    saver=tf.train.Saver(tf.global_variables())\n",
    "    with tf.Session() as sess:\n",
    "        #参数恢复\n",
    "        ckpt = tf.train.get_checkpoint_state(PATH)\n",
    "        if ckpt and ckpt.model_checkpoint_path:\n",
    "            saver.restore(sess, ckpt.model_checkpoint_path)\n",
    "            print(\"Load!\")\n",
    "        else:\n",
    "            tf.global_variables_initializer().run()\n",
    "        # module_file = tf.train.latest_checkpoint()\n",
    "        # saver.restore(sess, module_file)\n",
    "        test_predict=[]\n",
    "        for step in range(len(test_x)-1):\n",
    "          prob=sess.run(pred,feed_dict={X:[test_x[step]]})   \n",
    "          predict=prob.reshape((-1))\n",
    "          test_predict.extend(predict)\n",
    "        test_y=np.array(test_y)*std[7]+mean[7]\n",
    "        test_predict=np.array(test_predict)*std[7]+mean[7]\n",
    "        acc=np.average(np.abs(test_predict-test_y[:len(test_predict)])/test_y[:len(test_predict)])  #偏差\n",
    "        #以折线图表示结果\n",
    "        plt.figure()\n",
    "        plt.plot(list(range(len(test_predict))), test_predict, color='b')\n",
    "        plt.plot(list(range(len(test_y))), test_y,  color='r')\n",
    "        plt.show()\n",
    "\n",
    "prediction() "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
